package com.cg.ai.service.impl;

import com.cg.ai.entity.po.ChatMessage;
import com.cg.ai.entity.po.ChatSession;
import com.cg.ai.entity.vo.ChatEntity;
import com.cg.ai.entity.vo.ChatResponse;
import com.cg.ai.enums.SSEMsgType;
import com.cg.ai.mapper.ChatMessageMapper;
import com.cg.ai.mapper.ChatSessionMapper;
import com.cg.ai.service.ChatService;
import com.cg.ai.utils.SSEServer;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import reactor.core.publisher.Flux;

import java.time.Instant;
import java.util.List;

@Service
@RequiredArgsConstructor
@Slf4j
public class ChatServiceImpl implements ChatService {

    private final ChatSessionMapper chatSessionMapper;
    private final ChatMessageMapper chatMessageMapper;
    private final ChatClient chatClient;

    /**
     * 若传入的会话ID存在则直接返回该会话，否则创建一个默认的新会话。
     */
    @Override
    @Transactional
    public ChatSession createSessionIfNeeded(Long userId, Long sessionIdHint) {
        if (sessionIdHint != null) {
            ChatSession existingSession = chatSessionMapper.selectById(sessionIdHint);
            if (existingSession != null) {
                return existingSession;
            }
        }
        return newSession(userId);
    }

    @Override
    @Transactional
    public ChatSession createSession(Long userId, String title) {
        ChatSession session = newSession(userId);
        if (title != null && !title.isBlank()) {
            session.setSessionTitle(title);
            chatSessionMapper.update(session);
        }
        return session;
    }

    /**
     * 创建一个默认属性的新会话
     */
    private ChatSession newSession(Long userId) {
        ChatSession session = new ChatSession();
        session.setUserId(userId);
        session.setSessionTitle("新对话");
        session.setStatus("active");
        session.setMessageCount(0);
        Instant now = Instant.now();
        session.setCreateTime(now);
        session.setUpdateTime(now);
        session.setLastActiveTime(now);
        
        chatSessionMapper.insert(session);
        return session;
    }

    @Override
    @Transactional
    public ChatMessage appendUserMessage(Long sessionId, Long userId, String content) {
        return appendMessage(sessionId, userId, "user", content, "text");
    }

    @Override
    @Transactional
    public ChatMessage appendAssistantMessage(Long sessionId, Long userId, String content) {
        return appendMessage(sessionId, userId, "assistant", content, "text");
    }

    /**
     * 统一的消息追加逻辑
     */
    private ChatMessage appendMessage(Long sessionId, Long userId, String role, String content, String type) {
        int count = chatMessageMapper.countBySessionId(sessionId);
        ChatMessage message = new ChatMessage();
        message.setSessionId(sessionId);
        message.setUserId(userId);
        message.setRole(role);
        message.setContent(content);
        message.setMessageType(type);
        message.setMessageOrder(count + 1);
        Instant now = Instant.now();
        message.setCreateTime(now);
        message.setUpdateTime(now);
        
        chatMessageMapper.insert(message);
        
        // 更新会话的消息计数和最后活动时间
        chatSessionMapper.updateSessionActivity(sessionId, count + 1);
        
        return message;
    }

    @Override
    public List<ChatMessage> listMessages(Long sessionId) {
        return chatMessageMapper.selectBySessionId(sessionId);
    }

    @Override
    public List<ChatSession> listSessions(Long userId) {
        return chatSessionMapper.selectByUserId(userId);
    }

    @Override
    @Transactional
    public void deleteSession(Long sessionId) {
        // 首先删除会话下的所有消息
        chatMessageMapper.deleteBySessionId(sessionId);
        // 然后删除会话
        chatSessionMapper.deleteById(sessionId);
    }

    @Override
    @Transactional
    public ChatResponse chat(Long userId, Long sessionIdHint, String prompt) {
        ChatSession session = createSessionIfNeeded(userId, sessionIdHint);
        Long sid = session.getSessionId();
        appendUserMessage(sid, userId, prompt);

        List<Message> msgList = new java.util.ArrayList<>();
        List<ChatMessage> history = listMessages(sid);
        int start = Math.max(0, history.size() - 10);
        for (int i = start; i < history.size(); i++) {
            ChatMessage m = history.get(i);
            if (m == null) continue;
            String role = m.getRole();
            String content = m.getContent();
            if (content == null || content.isBlank()) continue;
            if ("user".equalsIgnoreCase(role)) msgList.add(new UserMessage(content));
            else if ("assistant".equalsIgnoreCase(role)) msgList.add(new AssistantMessage(content));
            else if ("system".equalsIgnoreCase(role)) msgList.add(new SystemMessage(content));
        }

        String answer = chatClient
                .prompt()
                .messages(msgList)
                .call()
                .content();

        appendAssistantMessage(sid, userId, answer);

        ChatResponse resp = new ChatResponse();
        resp.setSessionId(sid);
        resp.setAnswer(answer);
        return resp;
    }
    
    @Override
    public void doChat(ChatEntity chatEntity) {
        String sseUserKey = chatEntity.getCurrentUserName();
        String prompt = chatEntity.getMessage();
        Long userId = chatEntity.getUserId();
        Long sessionIdHint = chatEntity.getSessionId();

        ChatSession session = createSessionIfNeeded(userId, sessionIdHint);
        Long sid = session.getSessionId();

        // 先保存本次用户消息
        appendUserMessage(sid, userId, prompt);

        // 载入最近历史（最多10条）
        List<Message> msgList = new java.util.ArrayList<>();
        List<ChatMessage> history = listMessages(sid);
        int start = Math.max(0, history.size() - 10);
        for (int i = start; i < history.size(); i++) {
            ChatMessage m = history.get(i);
            if (m == null) continue;
            String role = m.getRole();
            String content = m.getContent();
            if (content == null || content.isBlank()) continue;
            if ("user".equalsIgnoreCase(role)) msgList.add(new UserMessage(content));
            else if ("assistant".equalsIgnoreCase(role)) msgList.add(new AssistantMessage(content));
            else if ("system".equalsIgnoreCase(role)) msgList.add(new SystemMessage(content));
        }

        StringBuilder answerBuffer = new StringBuilder();

        Flux<String> stringFlux = chatClient
                .prompt()
                .messages(msgList)
                .stream()
                .content();

        stringFlux
                .doOnError(throwable -> {
                    log.error("AI Stream error: " + throwable.getMessage());
                    SSEServer.sendMsg(sseUserKey, "AI service error", SSEMsgType.FINISH);
                    SSEServer.close(sseUserKey);
                })
                .subscribe(
                        content -> {
                            answerBuffer.append(content);
                            SSEServer.sendMsg(sseUserKey, content, SSEMsgType.ADD);
                        },
                        error -> log.error("Error processing stream: " + error.getMessage()),
                        () -> {
                            String finalAnswer = answerBuffer.toString();
                            appendAssistantMessage(sid, userId, finalAnswer);
                            SSEServer.sendMsg(sseUserKey, "done", SSEMsgType.FINISH);
                            SSEServer.close(sseUserKey);
                        }
                );
    }
}